import os
import os.path as osp
import pickle

import matplotlib.pyplot as plt
import metis
import networkx as nx
import numpy as np
import scipy.sparse as sp
import torch

import lzyutil
from config import args
import file_cache
import logging

dynamic_datasets = {'dgraph-nearby', 'dgraph-fin'}

directed_graph = {'dgraph-nearby', 'dgraph-fin', 'stanford'}


class GraphDataset:
    def __init__(self, dataset_name, use_cache=True, overwrite=True):
        """
        :param dataset_name: name of the dataset
        """
        self._edge_list = None
        self._adj_coo: sp.coo_matrix = None
        self._adj = None
        self._numnodes = None
        self._num_edges = None
        self._norm_laplacian_coo = None
        self._norm_laplacian = None
        self._nx_graph = None
        self._nl_nx_graph = None
        self.dataset_name = dataset_name

        # 动态图部分
        self._is_dynamic = False
        self._max_ts = 0
        self._cur_ts = 0
        self._cur_active_index = 0
        self._edge_time = None

        self._load_graph(dataset_name, use_cache, overwrite)

    def _load_graph(self, dataset_name, use_cache, overwrite):
        cache_path = file_cache.to_cache_path(dataset_name)
        if use_cache and osp.exists(cache_path):
            with open(cache_path, 'rb') as fin:
                self._is_dynamic, self._edge_list, self._edge_time, self._max_ts = pickle.load(fin)

        elif dataset_name not in dynamic_datasets:
            # graph with format of .txt
            if osp.exists(f'./dataset/{dataset_name}.txt'):
                with open(f'./dataset/{dataset_name}.txt') as f:
                    edges = []

                    for line in f.readlines():
                        strs = line.split()[:2]
                        src, dest = int(strs[0]), int(strs[1])
                        edges.append((src, dest))
                        if dataset_name not in directed_graph:
                            edges.append((dest, src))
                    self._edge_list = np.array(edges)
            # .gml suffix
            elif osp.exists(f'./dataset/{dataset_name}.gml'):
                graph = nx.read_gml(f'./dataset/{dataset_name}.gml', label=None)

                node_to_id = {node: i for i, node in enumerate(graph.nodes())}

                self._edge_list = np.array([[node_to_id[edge[0]], node_to_id[edge[1]]] for edge in graph.edges()])
                if dataset_name not in directed_graph:
                    self._edge_list = np.concatenate(
                        (self._edge_list,
                         np.array([[node_to_id[edge[1]], node_to_id[edge[0]]] for edge in graph.edges()])),
                        axis=0)


        else:
            self._is_dynamic = args.use_dynamic_graph
            dynamic_dir = f"./dataset/dynamic/{dataset_name}"
            edge_index = np.load(f'{dynamic_dir}/edge_index.npy')
            edge_time = np.load(f'{dynamic_dir}/edge_time.npy')
            edge_time = edge_time - min(edge_time)
            self._max_ts = max(edge_time)  # 用于将时间点剪切成snapshot
            edge_active_index = np.argsort(edge_time)
            self._edge_time, self._edge_list = edge_time[edge_active_index], edge_index[edge_active_index]

        # calculate the size of the adjacent matrix
        adj_size = max(max(self._edge_list[:, 0]), max(self._edge_list[:, 1])) + 1
        print(f'edge list size: {self._edge_list.size}')

        if not self._is_dynamic:  # 动态图的coo_matrix根据时间戳实时构建
            self._adj_coo = sp.coo_matrix(
                (np.ones(len(self._edge_list)), (self._edge_list[:, 0], self._edge_list[:, 1])),
                shape=(adj_size, adj_size)
            )
            self._adj_csr = self._adj_coo.tocsr()
        self._numnodes = adj_size
        file_cache.register(dataset_name,
                            (self._is_dynamic, self._edge_list, self._edge_time, self._max_ts)
                            , overwrite=overwrite)

    def _generate_graph_feature(self):
        pass

    def get_edge_list(self, get_all=True):
        if get_all:
            return self._edge_list
        else:
            return self._edge_list[:self._cur_active_index]

    def get_edge_index(self, get_all=True) -> torch.Tensor:
        return torch.from_numpy(self.get_edge_list(get_all)).transpose(0, 1)

    def get_size(self):
        return self._numnodes

    def get_norm_laplacian_coo(self) -> sp.coo_matrix:
        """
        get the normalized graph Laplacian matrix of the graph dataset

        :return: coo_matrix form of the normalized Laplacian matrix
        """
        if self._norm_laplacian_coo is None or self._is_dynamic:
            self._norm_laplacian_coo = lzyutil.get_norm_laplacian_sparse(self.get_adjacency_coo(),
                                                                         enhance_self_loop=True)
        return self._norm_laplacian_coo

    def get_norm_laplacian(self) -> np.ndarray:
        if self._norm_laplacian is None or self._is_dynamic:
            self._norm_laplacian = self.get_norm_laplacian_coo().toarray()
        return self._norm_laplacian

    def get_adjacency_coo(self) -> sp.coo_matrix:
        if self._is_dynamic:
            return sp.coo_matrix((np.ones(self._cur_active_index),
                                  (self._edge_list[:self._cur_active_index, 0],
                                   self._edge_list[:self._cur_active_index, 1])),
                                 shape=(self._numnodes, self._numnodes))
        else:
            return self._adj_coo

    def get_adjacency_csr(self) -> sp.csr_matrix:
        if self._is_dynamic:
            return sp.csr_matrix((np.ones(self._cur_active_index),
                                  (self._edge_list[:self._cur_active_index, 0],
                                   self._edge_list[:self._cur_active_index, 1])),
                                 shape=(self._numnodes, self._numnodes))
        else:
            return self._adj_csr

    def get_adjacency(self) -> np.ndarray:
        return self.get_adjacency_coo().toarray()

    def get_nx_graph(self) -> nx.Graph:
        if self._is_dynamic:
            self._nx_graph = nx.from_edgelist(self._edge_list[:self._cur_active_index])
        elif self._nx_graph is None:
            self._nx_graph = nx.from_scipy_sparse_array(self.get_adjacency_coo())
        return self._nx_graph

    def get_sparse_iter(self) -> lzyutil.SparseGraphIter:
        return lzyutil.SparseGraphIter(self.get_adjacency_csr())

    @property
    def is_dynamic(self):
        return self._is_dynamic

    @property
    def max_ts(self):
        return self._max_ts if self._is_dynamic else 0

    @property
    def cur_ts(self):
        return self._cur_ts

    @cur_ts.setter
    def cur_ts(self, value: int):
        """
        修改当前ts时，自动加入或删除图中的边，
        :param value:
        :return:
        """
        if not self._is_dynamic:
            return
        if value == self._cur_ts:
            return
        if value > self._max_ts:
            value = self._max_ts
        new_ts_index = np.searchsorted(self._edge_time, value, side='right')

        self._cur_active_index = new_ts_index
        self._cur_ts = value


def test_dynamic_dataset():
    gd = GraphDataset(dataset_name='dgraph-nearby')
    gd.cur_ts = 500000
    nxg = gd.get_nx_graph()
    print(nxg.edges())
    isolated = list(nx.isolates(nxg))
    print(isolated)
    nxg.remove_nodes_from(list(isolated))
    nx.draw(nxg, with_labels=False, node_size=50)
    file_cache.save_figure("test.png")
    return gd


def test_static_dataset():
    gd = GraphDataset(dataset_name='jazz')
    G = gd.get_nx_graph()
    a, b = metis.part_graph(G, 3)
    nx.draw(gd.get_nx_graph(), with_labels=True, node_size=100)
    file_cache.save_figure("test_static.png")
    return gd

def test_find_col_sum():
    gd = GraphDataset(dataset_name='power')
    adj_coo = gd.get_adjacency_coo()
    print(np.asarray(adj_coo.sum(axis=0)).flatten().shape)


def test_find_adj():
    gd = GraphDataset(dataset_name='jazz')
    csr = gd.get_adjacency_csr()
    print(csr[[18, 19, 20]].nonzero())
    x1 = np.random.uniform(0, 1, 10)
    x2 = np.random.randint(0, 2, 10)
    print(x1, x2)
    cond1 = x1 > 0.2
    cond2 = x2 == 0
    print(cond1, cond2)
    print(cond1 & cond2)

    # np.where( < infect_rate, 1, 0)


if __name__ == '__main__':
    # test_dynamic_dataset()
    test_find_col_sum()

    # test_dynamic_dataset()
    # test_find_adj()
